//#include "point_cloud_normal.h"

#include "point_cloud_normal.h"
namespace Registration {
__global__ void kernel_compute_vertex_map(uint16_t *depth_image,
                                          float3 *point_cloud, int image_width,
                                          int image_height,
                                          const float depth_cutoff,
                                          const CameraParams camera_params) {
  int x = blockIdx.x * blockDim.x + threadIdx.x;
  int y = blockIdx.y * blockDim.y + threadIdx.y;
  if (x >= image_width || x < 0 || y >= image_height || y < 0)
    return;
  int current_index = y * image_width + x;

  float depth_value = static_cast<float>(depth_image[current_index]);

  depth_value = depth_value / camera_params.depth_scale;
  if (depth_value > depth_cutoff) {
    depth_value == 0;
  }

  float3 vertex;
  vertex.x =
      (x - camera_params.principal_x) * depth_value / camera_params.focal_x;
  vertex.y =
      (y - camera_params.principal_y) * depth_value / camera_params.focal_y;
  vertex.z = depth_value;
  point_cloud[current_index] = vertex;
}

__global__ void kernel_generate_depth_map(const float3 *input_point_cloud,
                                          int point_num, float *depth_data,
                                          int image_width, int image_height,
                                          const CameraParams &camera_params) {

  int thread_num = blockDim.x * gridDim.x * blockDim.y * gridDim.y;
  int x = threadIdx.x + blockIdx.x * blockDim.x;
  int y = threadIdx.y + blockIdx.y * blockDim.y;
  int offset = y * blockDim.x * gridDim.x + x;
  for (int idx = offset; idx < point_num; idx = idx + thread_num) {
    float point_x, point_y, point_z;
    point_x = input_point_cloud[idx].x;
    point_y = input_point_cloud[idx].y;
    point_z = input_point_cloud[idx].z;
    int u = Round((point_x * camera_params.focal_x) / point_z +
                  camera_params.principal_x);
    int v = Round((point_y * camera_params.focal_y) / point_y +
                  camera_params.principal_y);
    if (u < image_width && v < image_height) {
      depth_data[image_height * v + u] = point_z;
    }
  }
}

__global__ void kernel_compute_normal_map(float3 *point_cloud,
                                          float3 *point_normal, int image_width,
                                          int image_height) {
  int x = blockIdx.x * blockDim.x + threadIdx.x;
  int y = blockIdx.y * blockDim.y + threadIdx.y;

  if (x < 1 || x >= image_width - 1 || y < 1 || y >= image_height - 1)
    return;

  int left_index = y * image_width + x - 1;
  float3 left_vertex = point_cloud[left_index];
  int right_index = y * image_width + x + 1;
  float3 right_vertex = point_cloud[right_index];
  int upper_index = (y - 1) * image_width + x;
  float3 upper_vertex = point_cloud[upper_index];
  int lower_index = (y + 1) * image_width + x;
  float3 lower_vertex = point_cloud[lower_index];
  float3 normal;
  float3 current_point = point_cloud[y * image_width + x];

  if (current_point.z == 0 || left_vertex.z == 0 || right_vertex.z == 0 ||
      upper_vertex.z == 0 || lower_vertex.z == 0)
    normal = make_float3(0, 0, 0);
  else {
    float3 hor = make_float3(left_vertex.x - right_vertex.x,
                             left_vertex.y - right_vertex.y,
                             left_vertex.z - right_vertex.z);
    float3 ver = make_float3(upper_vertex.x - lower_vertex.x,
                             upper_vertex.y - lower_vertex.y,
                             upper_vertex.z - lower_vertex.z);
    normal = cross(hor, ver);
  }
  int current_index = y * image_width + x;
  norm(normal);
  point_normal[current_index] = normal;
}

__device__ float3 cross(float3 &vec_1, float3 &vec_2) {
  float3 cross_result;
  cross_result.x = vec_1.y * vec_2.z - vec_2.y * vec_1.z;
  cross_result.y = vec_2.x * vec_1.z - vec_1.x * vec_2.z;
  cross_result.z = vec_1.x * vec_2.y - vec_2.x * vec_1.y;
  return cross_result;
}
__device__ int Round(float target_data) {
  return ((target_data - static_cast<int>(target_data)) > 0.5
              ? (static_cast<int>(target_data) + 1)
              : (static_cast<int>(target_data)));
}
void compute_vertex_map(uint16_t *depth_map, float3 *vertex_map,
                        int image_width, int image_height,
                        const float depth_cutoff,
                        const CameraParams camera_params) {
  dim3 threads(32, 32);
  dim3 blocks((image_width + threads.x - 1) / threads.x,
              (image_height + threads.y - 1) / threads.y);
  kernel_compute_vertex_map<<<blocks, threads>>>(depth_map, vertex_map,
                                                 image_width, image_height,
                                                 depth_cutoff, camera_params);
  cudaThreadSynchronize();
}

void compute_normal_map(float3 *vertex_map, float3 *normal_map, int image_width,
                        int image_height) {
  dim3 threads(32, 32);
  dim3 blocks((image_width + threads.x - 1) / threads.x,
              (image_height + threads.y - 1) / threads.y);
  kernel_compute_normal_map<<<blocks, threads>>>(vertex_map, normal_map,
                                                 image_width, image_height);
  cudaThreadSynchronize();
}

__device__ void norm(float3 &input) {
  float x = input.x;
  float y = input.y;
  float z = input.z;

  float length = sqrt(x * x + y * y + z * z);

  input.x = x / length;
  input.y = y / length;
  input.z = z / length;
}

/*
void test(uint16_t*input, uint16_t*output, int image_width, int image_height)
{

        dim3 threads(32, 32);
        dim3 blocks((image_width + threads.x - 1) / threads.x, (image_height +
threads.y - 1) / threads.y); kernel_compute_normal_map << <blocks, threads >>
>(input, output, image_width, image_height); cudaThreadSynchronize();
}*/

} // namespace Registration